//
// Created by yellowdo on 2024/11/25.
//

#ifndef CUDNN_SPATIALTF_IMPL_HPP
#define CUDNN_SPATIALTF_IMPL_HPP

#include "cudnn.h"

#include "gpu/nvidia/sycl_cuda_utils.hpp"


namespace dnnl {
namespace impl {
namespace gpu {
namespace nvidia {

struct cudnn_spatialtf_impl_base_t {
    cudnnDataType_t data_type;
    int ndims;

    virtual ~cudnn_spatialtf_impl_base_t() {}

    virtual status_t init(const spatialtf_pd_t *pd) = 0;

    virtual void execute(cudnnHandle_t handle, void **x, int size) const = 0;


};

struct cudnn_spatialtf_gridgen_fwd_impl_t : public cudnn_spatialtf_impl_base_t {

    cudnnSpatialTransformerDescriptor_t spatial_tf_desc;
    int dims[DNNL_MAX_NDIMS];
    cudnnDataType_t data_type;

    status_t init(const spatialtf_pd_t *pd) override {
        convert_data_type(pd->src_md(), &data_type);
        int ndims = pd->src0_ndims();
        pd->get_src0_dims(dims);

        create_and_set_spatial_tf_descriptor(&spatial_tf_desc, data_type, ndims, dims);
        return status::success;
     }


    void execute(cudnnHandle_t handle, void **x, int size) const override {
        std::cout << "cudnn_spatialtf_gridgen_fwd_impl_t execute" << std::endl;
        CUDNN_EXECUTE_FUNC(cudnnSpatialTfGridGeneratorForward, handle,spatial_tf_desc, x[0], x[1]);

    }

    ~cudnn_spatialtf_gridgen_fwd_impl_t() {
        CUDNN_EXECUTE_FUNC_V(cudnnDestroySpatialTransformerDescriptor, spatial_tf_desc);
    }
};

} // namespace nvidia
} // namespace gpu
} // namespace impl
} // namespace dnnl

#endif //CUDNN_SPATIALTF_IMPL_HPP
